import albumentations as A
import cv2
from PIL import Image
# from torchvision import transforms
import random
import os


def create_newimg(trans, img):
    assert (trans != None)

    return trans(image=img)["image"]


def aug_face(srcroot, dstroot):
    light_trans = A.Compose([

        A.RandomBrightnessContrast(p=1),  # 随机亮度调整

    ])
    motion_trans = A.Compose([

        A.MotionBlur(p=1),  # 动态模糊

    ])

    for folder in os.listdir(srcroot):
        load_folder = fr"{srcroot}/{folder}"
        sa_folder = fr"{dstroot}/{folder}"
        if not os.path.exists(sa_folder):
            os.mkdir(sa_folder)
        #采用了动态模糊和亮度调节
        for index, img_name in enumerate(os.listdir(load_folder)):
            img0 = cv2.imread(fr"{load_folder}/{img_name}")
            img1 = create_newimg(light_trans, img0)
            img2 = create_newimg(motion_trans, img0)
            cv2.imwrite(fr"{sa_folder}/{index * 3}.png", img0, (cv2.IMWRITE_PNG_COMPRESSION, 100))
            cv2.imwrite(fr"{sa_folder}/{index * 3 + 1}.png", img1, (cv2.IMWRITE_PNG_COMPRESSION, 100))
            cv2.imwrite(fr"{sa_folder}/{index * 3 + 2}.png", img2, (cv2.IMWRITE_PNG_COMPRESSION, 100))
# 生成测试txt文件 默认相同人500，不同人500
def gen_paris_txt(root,paris=1000):

    assert (int(paris%2) == 0)

    # 1.获取文件名list
    img_list = []
    same_list = set()
    diff_list = set()

    for folder in os.listdir(root):

        for file_name in os.listdir(fr"{root}/{folder}"):
            item = fr"{folder}/{file_name}"
            if item not in img_list:
                img_list.append(item)



    # 2.随机取出两个文件名，如果名字相同，则加入到same_list,否则加入到diff_list

    img_length = len(img_list)
    # print(img_length)
    while True:
        # 采样完成，推出循环
        if len(same_list) >= paris//2 and len(diff_list) >= paris//2:
            break
        i = random.randint(0,img_length-1)
        j = random.randint(0,img_length-1)
        print('----------------')
        print(i,j)
        file1 = img_list[i].split(r'/')[0]
        file2 = img_list[j].split(r'/')[0]
        print(file1,file2)
        print(len(same_list),len(diff_list))
        # same list 判断
        if file1 == file2:
            if len(same_list) >= 500:
                continue
            if (i,j) not in same_list and (j,i) not in same_list:
                same_list.add((i,j))

        # diff list 判断
        if file1 != file2:
            if len(diff_list) >= 500:
                continue
            if (i, j) not in diff_list and (j, i) not in diff_list:
                diff_list.add((i, j))



    # 3.将list写入txt文件
    with open('lw_pairs.txt','w') as f:
        for (i,j) in same_list:
            f.write(f'{img_list[i]} {img_list[j]} 1\n')
        for (i,j) in diff_list:
            f.write(f'{img_list[i]} {img_list[j]} 0\n')



if __name__ == '__main__':
    path = r'c:\Users\2\Desktop\1.jpg'
    # aug_face(r"c:\Users\2\Desktop\19studentFace", r"c:\Users\2\Desktop\test")
    gen_paris_txt(r"c:\Users\2\Desktop\test")